import torch
from torch import nn

"""
用 PyTorch 原生的 nn.Transformer / nn.TransformerEncoder / nn.TransformerDecoder 时，Embedding 和位置编码都不自带，要自己写（或用别的库的封装）。
这些模块只实现了“注意力 + FFN + Add&Norm”等层的堆叠；不包含：
token→向量的 nn.Embedding
位置编码（正弦、可学习、相对位置/RoPE 等）
词表线性输出（Linear(d_model → vocab_size)）
只有一些高层封装/预训练模型（如 HuggingFace 的 BertModel/GPT2LMHeadModel、torchvision 的 ViT 系列）会把嵌入 + 位置信息集成在内部。

需要自己写：nn.Embedding、位置编码（或相对位置/旋转位置等）、输出 Linear、各类 mask（padding/因果）。
不需要自己写：多头注意力、前馈网络、残差 + LayerNorm，这些 nn.Transformer* 已经提供。
如果用预训练模型（HuggingFace 等）：通常自带词嵌入和位置编码，你只管喂 token id。
"""

"""
小提示：d_model 最好为偶数；若是奇数，这个实现最后会少一个奇数通道的 cos 列（通常 Transformer 设为偶数）。
"""


class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        """
        d_model: 模型的特征维度
        max_len: 序列最大长度，默认5000
        dropout: Dropout比率，默认0.1
        """
        super(PositionalEncoding, self).__init__()
        # 创建一个全零张量pe，形状为[max_len, d_model]，用于存储位置编码
        pe = torch.zeros(max_len, d_model)
        # 生成位置索引序列[0,1,2,...,max_len-1]，并增加一个维度变为列向量[max_len, 1]
        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # 计算位置编码的分母项，用于正弦和余弦函数的频率调节
        # torch.arange(0, d_model, 2)生成偶数索引[0,2,4,...]
        # -torch.log(torch.tensor(10000.0)) / d_model是缩放因子
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
        # 对偶数维度(0,2,4,...)应用正弦函数计算位置编码
        pe[:, 0::2] = torch.sin(pos * div)
        # 对奇数维度(1,3,5,...)应用余弦函数计算位置编码
        pe[:, 1::2] = torch.cos(pos * div)
        # 将位置编码注册为缓冲区，不参与梯度更新
        # 增加批次维度，形状变为[1, max_len, d_model]
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        # 创建Dropout层
        # nn.Dropout(dropout) 创建一个丢弃层，用于防止过拟合
        # self.dropout 将该丢弃层保存为类的属性，以便在前向传播时使用
        # dropout参数控制随机将输入张量中一定比例的元素置零的概率
        self.dropout = nn.Dropout(dropout)

    """
    广播机制工作原理
        x 的维度：[batch_size, seq_len, d_model]
        self.pe[:, :L, :] 的维度：[1, L, d_model]
    相加过程
        广播扩展：由于 self.pe[:, :L, :] 的第一个维度是1，它会通过广播机制自动扩展到与 x 相同的 batch_size
        逐元素相加：每个批次的数据都会与相同的位置编码进行相加
    这意味着同一个位置编码会被应用到批次中的每一个样本，这正是位置编码的设计目的——相同位置的词应该具有相同的位置信息。
    """

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        在 PositionalEncoding 的 forward 方法中，参数 x 的维度是三维的，具体为：
            形状: [batch_size, seq_len, d_model]
                batch_size: 批次大小，表示同时处理的序列数量
                seq_len: 序列长度，表示每个序列中的元素个数
                d_model: 模型的特征维度，与位置编码的维度保持一致
        在前向传播过程中：
            通过 x.size(1) 获取序列长度 L
            从预计算的位置编码 self.pe 中切片出相应长度的位置编码 self.pe[:, :L, :]
            将位置编码加到输入张量 x 上，实现位置信息的注入
        这种三维设计使得位置编码可以同时处理一个批次中的多个序列，并且能够适应不同长度的序列输入。
        """
        # 获取输入序列长度
        L = x.size(1)
        # 将输入张量与位置编码相加，self.pe是预定义的位置编码矩阵
        # 将三维切片结果与输入张量 x 进行广播相加
        x = x + self.pe[:, :L, :]  # 切片操作
        # 对结果应用dropout操作并返回
        return self.dropout(x)


if __name__ == '__main__':
    PositionalEncoding(4, 5)
